//
// Created by whitby on 8/4/22.
//

#ifndef DENSESURFELMAPPING_SURFEL_MAP_H
#define DENSESURFELMAPPING_SURFEL_MAP_H

#include <list>
#include <vector>
#include <set>
#include <iostream>
#include <chrono>
#include <thread>
#include <Eigen/Eigen>
#include <opencv2/opencv.hpp>
#include <string>

#include "elements.h"
#include "fusion_functions.h"

using namespace std;

struct PoseElement {
public:
    vector<shared_ptr<SurfelElement>> attached_surfels;
    Eigen::Matrix4f *cam_pose = nullptr;
    Eigen::Matrix4f *loop_pose = nullptr;
    vector<int> linked_pose_index;

    PoseElement() = default;

    bool loopChanged() const {
        if (cam_pose == nullptr || loop_pose == nullptr)
            return false;

        if ((*cam_pose) != (*loop_pose))
            return true;
        else return false;
    }

    Eigen::Matrix4f getWarpPose() const {
        if (cam_pose == nullptr || loop_pose == nullptr) {
            return Eigen::Matrix4f::Identity();
        }
        return (*loop_pose) * (*cam_pose).inverse();
    }
};

class SurfelMap {
public:
    SurfelMap(int width, int height, Eigen::Matrix3f K, float far_dist, float near_dist, int drift_free_poses);

    ~SurfelMap();

    void save_cloud(const string &save_path);

    void orb_result_process(Eigen::Matrix4f &this_pose, int reference_index, bool is_new_keyFrame,
                            vector<Eigen::Matrix4f> &allKeyFrames, vector<pair<int, int>> &loop_path);

    void synchronize(cv::Mat &left_img, cv::Mat &depth_img, int reference_index, Eigen::Matrix4f &current_pose);

private:
    void fuse_map(cv::Mat &image, cv::Mat &depth, Eigen::Matrix4f &pose, int reference_index);

    void move_add_surfels(int reference_index);

    void get_add_remove_poses(int root_index, vector<int> &pose_to_add, vector<int> &pose_to_remove);

    void get_drift_free_poses(int root_index, vector<int> &driftFree_poses, int drift_free_range);

    // camera param
    int cam_width, cam_height;
    float cam_fx, cam_fy, cam_cx, cam_cy;
    Eigen::Matrix3f camera_matrix;

    // fusion tools
    FusionFunctions fusion_functions;

    // database
    vector<shared_ptr<SurfelElement>> local_surfels;
    vector<shared_ptr<SurfelElement>> inactivate_surfels;
    vector<PoseElement> poses_database;
    std::set<int> local_surfels_indexs;
    int drift_free_poses;

    std::vector<std::thread> warp_thread_pool;
    int warp_thread_num = 1;

    void warp_surfels();

    void warp_inactive_surfels_cpu_kernel(int thread_i, int step);

    void warp_activate_surfels_cpu_kernel(int thread_i, int thread_num, const Eigen::Matrix4f &transform_m);
};

#endif //DENSESURFELMAPPING_SURFEL_MAP_H
